#include "scan_matching/optimizer.h"

Optimizer::Optimizer(std::shared_ptr<OccupiedGridCostFunction2D> grid_cost_function_ptr) {
    grid_cost_function_ptr_ = grid_cost_function_ptr;
    robust_kernel_ = std::make_shared<HuberKernel>(1.0);
}

void Optimizer::Optimize(const int& max_iterations) {
    double mu = 1e-4;
    double v = 2.0;
    double min_error_reduction = 1e-6;
    double current_error = ComputeTotalError();
    auto init_vertices = vertices_;
    // cout << "current_error ==== : " << current_error << std::endl;
    is_solve_success_ = false;
    for(int iter = 0; iter < max_iterations; ++iter) {
       BuildSystem();
       auto backup_vertices = vertices_;

       // cout << "iter ==== : " << iter << std::endl;

       // LM阻尼
       MatrixXd H_lm = H_;
       for (int i = 0; i < H_lm.rows(); ++i) {
          H_lm(i, i) += mu * H_lm(i, i);
       }
       
       // cout << "H_lm ==== : " << H_lm << std::endl;
       // cout << "b_ ==== : " << b_ << std::endl;

       // 求解
       VectorXd dx;
       LDLT<MatrixXd> ldlt(H_lm);
       if (ldlt.info() == Success) {
          dx = ldlt.solve(-b_);
       } else {
          cerr << "Warning: LDLT failed, increasing damping." << endl;
          mu *= v;
          v *= 2.0;
          continue;
       }
    
       cout << "iter : " << iter << ", dx ==== : " << dx.transpose() << std::endl;

       UpdatePoses(dx);
       double new_error = ComputeTotalError();
       double error_reduction = current_error - new_error;

       // cout << "error_reduction ==== : " << error_reduction << std::endl;

       double predicted_reduction = 0.5 * (dx.dot(mu * H_.diagonal().cwiseProduct(dx))) - 0.5 * b_.dot(dx);

       double rho = (predicted_reduction > 0) ? (error_reduction / predicted_reduction) : 0;

       // cout << "rho : " << rho << std::endl;

       if (rho >= 0.001) { 
           current_error = new_error;
           mu *= max(1.0/3.0, 1.0 - pow(2.0*rho - 1.0, 3));
           v = 2.0;

           if (abs(error_reduction) < min_error_reduction) {
              is_solve_success_ = true;

              cout << "is_solve_success_ ==== : " << is_solve_success_ << std::endl;

              break;
           }
       } else {
           vertices_ = backup_vertices;
           mu *= v;
           v *= 2.0;
       }
    }

    if(!is_solve_success_) {
        vertices_ = init_vertices;
    }
}

void Optimizer::UpdatePoses(const VectorXd& dx) {
    for (const auto& v : var_indices_) {
        int idx = v.second;
        vertices_[v.first].p[0] += dx(3*idx);
        vertices_[v.first].p[1] += dx(3*idx + 1);
        vertices_[v.first].R = vertices_[v.first].R * SO2(dx(3*idx + 2));
    }
}

int Optimizer::AddVertex(const State& state, const bool& is_fixed) {
    int id = next_vertex_id_++;
    vertices_[id] = state;
    is_fixed_[id] = is_fixed;
    vertex_queue_.push(id);

    // cout << "id ==== : " << id << std::endl;
    // cout << "vertex_queue_.size() ==== : " << vertex_queue_.size() << std::endl;

    if (vertex_queue_.size() > window_size_) {
        MarginalizeOldestVertex();
    }

    return id;
}

void Optimizer::MarginalizeOldestVertex() {
    if (vertex_queue_.empty()) return;
        int marg_id = vertex_queue_.front();
        vertex_queue_.pop();

        unordered_set<size_t> related_edges = vertex_to_edges_[marg_id];

        cout << "marg_id ==== : " << marg_id << std::endl;

        // cout << "related_edges ==== : " << vertex_to_edges_.size() << std::endl;
        // cout << "vertex_queue_ ==== : " << vertex_queue_.size() << std::endl;

        unordered_set<int> related_vertices;
        for (size_t edge_idx : related_edges) {
            const auto& edge = edges_[edge_idx];

            if (edge.id1 != marg_id && vertices_.count(edge.id1)) {
                related_vertices.insert(edge.id1);
            }
            if (edge.id2 != marg_id && vertices_.count(edge.id2)) {
                related_vertices.insert(edge.id2);
            }
        }

        // 构建局部系统
        unordered_map<int, int> local_indices;
        int idx = 0;
        local_indices[marg_id] = idx++;
        for (int vid : related_vertices) {
            local_indices[vid] = idx++;
        }

        MatrixXd H_local = MatrixXd::Zero(3*idx, 3*idx);
        VectorXd b_local = VectorXd::Zero(3*idx);

        // cout << "related_edges ==== : " << related_edges.size() << std::endl;

        // 填充局部H和b
        for (size_t edge_idx : related_edges) {
            const auto& edge = edges_[edge_idx];
            if (!vertices_.count(edge.id1) || !vertices_.count(edge.id2)) continue;
            
            State xi = vertices_.at(edge.id1), xj = vertices_.at(edge.id2);

            // cout << "edge.id1 ==== : " << edge.id1 << ", " << edge.id2 << std::endl;
            
            int idx1 = local_indices[edge.id1], idx2 = local_indices[edge.id2];

            if(edge.measurement.type == "g") {
                Eigen::Matrix<double, 1, 1> e;
                Eigen::Matrix<double, 1, 3> jacobianXi;
            
                // cout << "grid_cost_function_ptr_->GetPointCloud().size() ==== : " << grid_cost_function_ptr_->GetPointCloud().size() << std::endl;

                for (size_t i = 0; i < grid_cost_function_ptr_->GetPointCloud().size(); ++i) {
                    // Note that this is a 2D point. The third component is a scaling factor.
                    Eigen::Matrix<double, 3, 1> point(((grid_cost_function_ptr_->GetPointCloud()[i].position.x())),
                                            ((grid_cost_function_ptr_->GetPointCloud()[i].position.y())),
                                            (1.));
                    if(!grid_cost_function_ptr_->ComputeResidualAndJacobianPoseGraph(xi, point, e, jacobianXi)) {
                        continue;
                    }

                    MatrixXd weighted_H;
                    VectorXd weighted_e;
                    std::tie(weighted_H, weighted_e) = robust_kernel_->robustify(e);
            
                    Eigen::MatrixXd Omega_robust = edge.information * weighted_H;
                    VectorXd e_robust = weighted_e;    

                    if (!is_fixed_[edge.id1]) {
                        H_.block<3, 3>(3*idx1, 3*idx1) += jacobianXi.transpose() * Omega_robust * jacobianXi;
                        b_.segment<3>(3*idx1) += jacobianXi.transpose() * Omega_robust * e_robust;
                    } 
                }
            }

            // cout << "edge.id1 ==== : " << edge.id1 << std::endl;
            // cout << "edge.id2 ==== : " << edge.id2 << std::endl;
            
            if(edge.measurement.type == "p") {
                Vector3d e;
                Eigen::Matrix<double, 3, 3> jacobianXi, jacobianXj;

                if(edge.id1 == edge.id2) {
                   PoseObserve::ComputeResidualAndJacobianPoseGraph(xi, edge.measurement, e, jacobianXi);
                   MatrixXd weighted_H;
                   VectorXd weighted_e;
                   std::tie(weighted_H, weighted_e) = robust_kernel_->robustify(e);
            
                   Matrix3d Omega_robust = edge.information * weighted_H;
                   Vector3d e_robust = weighted_e;
                   // cout << "p idx1 ==== : " << idx1 << " , " << idx2 << std::endl;

                   H_local.block<3, 3>(3*idx1, 3*idx1) += jacobianXi.transpose() * Omega_robust * jacobianXi;
                   b_local.segment<3>(3*idx1) += jacobianXi.transpose() * Omega_robust * e_robust;
                }
            }    
        }

        // cout << "related_vertices ==== : " << related_vertices.size() << std::endl;

        // 舒尔补计算
        if (related_vertices.size() > 0) {
            MatrixXd H_aa = H_local.block<3,3>(0, 0);
            MatrixXd H_ab = H_local.block(0, 3, 3, 3*(idx-1));
            MatrixXd H_ba = H_local.block(3, 0, 3*(idx-1), 3);
            MatrixXd H_bb = H_local.block(3, 3, 3*(idx-1), 3*(idx-1));

            VectorXd b_a = b_local.segment<3>(0);
            VectorXd b_b = b_local.segment(3, 3*(idx-1));

            // 确保H_aa可逆
            SelfAdjointEigenSolver<Matrix3d> eigensolver(H_aa);
            if (eigensolver.eigenvalues().minCoeff() < 1e-6) {
                H_aa += Matrix3d::Identity() * 1e-4;
            }

            MatrixXd H_aa_inv = H_aa.ldlt().solve(Matrix3d::Identity());
            MatrixXd H_prior = H_bb - H_ba * H_aa_inv * H_ab;
            VectorXd b_prior = b_b - H_ba * H_aa_inv * b_a;

            // 存储为先验因子
            MarginalizationFactor factor;
            factor.H_prior = H_prior;
            factor.b_prior = b_prior;
            factor.related_vertices = related_vertices;
            marginalization_factors_[marg_id] = factor;
        }

        // 清理
        vertices_.erase(marg_id);
        is_fixed_.erase(marg_id);
        for (size_t edge_idx : related_edges) {
            if (active_edges_.count(edge_idx)) {
                const auto& edge = edges_[edge_idx];
                vertex_to_edges_[edge.id1].erase(edge_idx);
                vertex_to_edges_[edge.id2].erase(edge_idx);
                active_edges_.erase(edge_idx);
            }
        }
        vertex_to_edges_.erase(marg_id);
}

double Optimizer::ComputeTotalError() {
    double total_error = 0.0;
    for (const auto& edge_idx : active_edges_) {
        const auto& edge = edges_[edge_idx];
        if (!vertices_.count(edge.id1) || !vertices_.count(edge.id2)) continue;
            
        State xi = vertices_.at(edge.id1), xj = vertices_.at(edge.id2);

        if(edge.measurement.type == "g") {
            Eigen::Matrix<double, 1, 1> e;
            Eigen::Matrix<double, 1, 3> jacobianXi;
            for (size_t i = 0; i < grid_cost_function_ptr_->GetPointCloud().size(); ++i) {
                // Note that this is a 2D point. The third component is a scaling factor.
                Eigen::Matrix<double, 3, 1> point(((grid_cost_function_ptr_->GetPointCloud()[i].position.x())),
                                            ((grid_cost_function_ptr_->GetPointCloud()[i].position.y())),
                                            (1.));
                if(!grid_cost_function_ptr_->ComputeResidualAndJacobianPoseGraph(xi, point, e, jacobianXi)) {
                    continue;
                }
                total_error +=  e.transpose() * edge.information * e;
            }
        }

        if(edge.measurement.type == "p") {
            Vector3d e;
            Eigen::Matrix<double, 3, 3> jacobianXi, jacobianXj;
            if(edge.id1 == edge.id2) {
                PoseObserve::ComputeResidualAndJacobianPoseGraph(xi, edge.measurement, e, jacobianXi);
                total_error += e.transpose() * edge.information * e;
            }
        }    
    }
    return total_error;
}

void Optimizer::AddEdge(const PoseGraphEdge& edge) {
    edges_.push_back(edge);
    size_t edge_idx = edges_.size() - 1;
    active_edges_.insert(edge_idx);
    vertex_to_edges_[edge.id1].insert(edge_idx);
    vertex_to_edges_[edge.id2].insert(edge_idx);

    // cout << "edge ==== : " << edge.measurement.type << std::endl;

    // cout << "edge.id ==== : " << edge.id1 << " , " << edge.id2 << std::endl;
    // cout << "edge_idx ==== : " << edge_idx << std::endl;
}

void Optimizer::BuildSystem() {
    var_indices_.clear();
    int free_vars = 0;
    
    // cout << "vertices_ ==== : " << vertices_.size() << std::endl;

    for (const auto& v : vertices_) {
        if (!is_fixed_[v.first]) {
            var_indices_[v.first] = free_vars++;
        }
    }

    // cout << "free_vars ==== : " << free_vars << std::endl;

    H_ = MatrixXd::Zero(3 * free_vars, 3 * free_vars);
    b_ = VectorXd::Zero(3 * free_vars);

    // 添加边约束
    for (size_t i = 0; i < edges_.size(); ++i) {
        if (!active_edges_.count(i)) continue;
            
        const auto& edge = edges_[i];
        if (!vertices_.count(edge.id1) || !vertices_.count(edge.id2)) continue;
        
        State xi = vertices_.at(edge.id1), xj = vertices_.at(edge.id2);
        
        int idx1 = var_indices_[edge.id1], idx2 = var_indices_[edge.id2];

        // cout << "edge.id1 ==== : " << edge.id1 << ", " << edge.id2 << std::endl;

        if(edge.measurement.type == "g") {
            Eigen::Matrix<double, 1, 1> e;
            Eigen::Matrix<double, 1, 3> jacobianXi;
            
            // cout << "grid_cost_function_ptr_->GetPointCloud().size() ==== : " << grid_cost_function_ptr_->GetPointCloud().size() << std::endl;

            for (size_t i = 0; i < grid_cost_function_ptr_->GetPointCloud().size(); ++i) {
                // Note that this is a 2D point. The third component is a scaling factor.
                Eigen::Matrix<double, 3, 1> point(((grid_cost_function_ptr_->GetPointCloud()[i].position.x())),
                                            ((grid_cost_function_ptr_->GetPointCloud()[i].position.y())),
                                            (1.));
                if(!grid_cost_function_ptr_->ComputeResidualAndJacobianPoseGraph(xi, point, e, jacobianXi)) {
                    continue;
                }

                MatrixXd weighted_H;
                VectorXd weighted_e;
                std::tie(weighted_H, weighted_e) = robust_kernel_->robustify(e);
            
                Eigen::MatrixXd Omega_robust = edge.information * weighted_H;
                VectorXd e_robust = weighted_e;    

                // cout << "jacobianXi ==== : " << jacobianXi << std::endl;

                // cout << "Omega_robust ==== : " << Omega_robust << std::endl;
                // cout << "e ==== : " << e << std::endl;
                // cout << "e_robust ==== : " << e_robust << std::endl;

                if (!is_fixed_[edge.id1]) {
                    H_.block<3, 3>(3*idx1, 3*idx1) += jacobianXi.transpose() * Omega_robust * jacobianXi;
                    b_.segment<3>(3*idx1) += jacobianXi.transpose() * Omega_robust * e_robust;
                } 
            }
        }

        if(edge.measurement.type == "p") {
            Vector3d e;
            Eigen::Matrix<double, 3, 3> jacobianXi;

            if(edge.id1 == edge.id2) {
                PoseObserve::ComputeResidualAndJacobianPoseGraph(xi, edge.measurement, e, jacobianXi);
                MatrixXd weighted_H;
                VectorXd weighted_e;
                std::tie(weighted_H, weighted_e) = robust_kernel_->robustify(e);
            
                Matrix3d Omega_robust = edge.information * weighted_H;
                Vector3d e_robust = weighted_e;
                            
                // cout << "p idx1 ==== : " << idx1 << " , " << idx2 << std::endl;

                if (!is_fixed_[edge.id1] && !is_fixed_[edge.id2]) {
                    H_.block<3, 3>(3*idx1, 3*idx1) += jacobianXi.transpose() * Omega_robust * jacobianXi;
                    b_.segment<3>(3*idx1) += jacobianXi.transpose() * Omega_robust * e_robust;
                } else if (!is_fixed_[edge.id1]) {
                    H_.block<3, 3>(3*idx1, 3*idx1) += jacobianXi.transpose() * Omega_robust * jacobianXi;
                    b_.segment<3>(3*idx1) += jacobianXi.transpose() * Omega_robust * e_robust;
                } 
            }
        }
    }

    // 添加边缘化先验因子
    for (const auto& factor_pair : marginalization_factors_) {
        const auto& factor = factor_pair.second;
        vector<int> active_vertices;
            
        for (int vid : factor.related_vertices) {
            if (vertices_.count(vid) && !is_fixed_[vid]) {
                active_vertices.push_back(vid);
            }
        }

        if (active_vertices.empty()) continue;

        MatrixXd H_local = MatrixXd::Zero(3 * active_vertices.size(), 
                                          3 * active_vertices.size());
        VectorXd b_local = VectorXd::Zero(3 * active_vertices.size());

        for (size_t i = 0; i < active_vertices.size(); ++i) {
            int global_i = var_indices_[active_vertices[i]];
            for (size_t j = 0; j < active_vertices.size(); ++j) {
                int global_j = var_indices_[active_vertices[j]];
                H_local.block<3,3>(3*i, 3*j) = factor.H_prior.block<3,3>(3*global_i, 3*global_j);
            }
            b_local.segment<3>(3*i) = factor.b_prior.segment<3>(3*global_i);
        }

        for (size_t i = 0; i < active_vertices.size(); ++i) {
            int global_i = var_indices_[active_vertices[i]];
            for (size_t j = 0; j < active_vertices.size(); ++j) {
                int global_j = var_indices_[active_vertices[j]];
                H_.block<3,3>(3*global_i, 3*global_j) += H_local.block<3,3>(3*i, 3*j);
            }
            b_.segment<3>(3*global_i) += b_local.segment<3>(3*i);
        }
    }
}

void Optimizer::UpdateState(State& state, const int id) {
    auto iter = vertices_.find(id);
    if(iter != vertices_.end()) {
        state = iter->second;
    }
}
